#include "gifencoder.h"
#include "lzwencoder.h"
#include "quantization/grayscalequantizer.h"
#include "quantization/mediancutquantizer.h"
#include "quantization/mostusedquantizer.h"
#include "quantization/neuralquantizer.h"
#include "quantization/octreequantizer.h"
#include "quantization/palettequantizer.h"
#include <QDataStream>
#include <QtCore/qmath.h>

namespace WingGif {

GifEncoder::GifEncoder(QIODevice *stream, QObject *parent) : QObject(parent) {
  internalStream = stream;
}

GifEncoder::~GifEncoder() {
  // Add a comment section.
  writeComment("Made with WingGif");
  // Complete the file.
  QDataStream p(internalStream);
  p << uchar(0x3b);
  // Push data.
  internalStream->close();
}

int GifEncoder::repeatCount() { return _repeatCount; }

bool GifEncoder::setRepeatCount(int value) {
  if (value < 0)
    return false;

  _repeatCount = value;
  return true;
}

bool GifEncoder::useFullTransparency() { return _useFullTransparency; }

void GifEncoder::setUseFullTransparency(bool value) {
  _useFullTransparency = value;
}

QColor GifEncoder::transparentColor() { return _transparentColor; }

void GifEncoder::setTransparentColor(QColor value) {
  _transparentColor = value;
}

int GifEncoder::maximumNumberColor() { return _maximumNumberColor; }

bool GifEncoder::setMaximumNumberColor(int value) {
  if (value <= 0 || value > 256)
    return false;

  _maximumNumberColor = value;
  return true;
}

ColorQuantizationTypes GifEncoder::quantizationType() {
  return _quantizationType;
}

void GifEncoder::setQuantizationType(ColorQuantizationTypes value) {
  _quantizationType = value;
}

bool GifEncoder::useGlobalColorTable() { return _useGlobalColorTable; }

void GifEncoder::setUseGlobalColorTable(bool value) {
  _useGlobalColorTable = value;
}

int GifEncoder::samplingFactor() { return _samplingFactor; }

void GifEncoder::setSamplingFactor(int value) { _samplingFactor = value; }

void GifEncoder::addFrame(QByteArray pixels, QRect rect, int delay,
                          bool isLastFrame) {
  readPixels(pixels);

  // For global color table, only generate a new palette if it's the first
  // frame.
  if (!_useGlobalColorTable || isFirstFrame)
    calculateColorTableSize();

  if (isFirstFrame) {
    writeLogicalScreenDescriptor(rect);

    // Global color table.
    if (_useGlobalColorTable)
      writePalette();

    if (_repeatCount > -1)
      writeApplicationExtension();
  }

  writeGraphicControlExtension(delay, isLastFrame);
  writeImageDescriptor(rect);

  isFirstFrame = false;

  // Local color table.
  if (!_useGlobalColorTable)
    writePalette();

  writeImage();
}

void GifEncoder::writeLogicalScreenDescriptor(QRect rect) {
  // File Header, 6 bytes
  internalStream->write("GIF89a", 6);

  // Initial Logical Size (Width, Height), 4 bytes
  QDataStream p(internalStream);
  p << ushort(rect.width()) << ushort(rect.height());

  // Packed fields, 1 byte
  // Color resolution: 111 = (8 bits - 1)
  // Color depth - 1
  // Global colors count = 2^color depth
  // Sort flag (for the global color table): 0
  uchar pixelBits =
      uchar((colorTableSize & 7) << 1 | _useGlobalColorTable | 16);

  // Size of the Global Color Table (Zero, if not used.):
  uchar sizeInBits = _useGlobalColorTable ? (colorTableSize & 7) : 0;

  auto bitpak = pixelBits | (sizeInBits << 5);
  p << uchar(bitpak)
    << uchar(_useFullTransparency ? findTransparentColorIndex()
                                  : 0) // Background color index, 1 byte
    << uchar(0); // Pixel aspect ratio - Assume 1:1, 1 byte
}

void GifEncoder::writePalette() {
  QDataStream p(internalStream);

  for (auto color : colorTable) {
    p << uchar(color.red()) << uchar(color.green()) << uchar(color.blue());
  }

  // Do I need to fill up the rest of the color table?
  // Or just seek the stream to the next place?

  //(MaximumColorsCount -  ColorCount) * 3 channels [rgb]
  auto emptySpace = (getMaximumColorCount() - colorTable.count()) * 3;
  internalStream->write(QByteArray(emptySpace, char(0)));
}

void GifEncoder::writeApplicationExtension() {
  QDataStream p(internalStream);
  p << uchar(0x21)  // Extension Introducer.
    << uchar(0xff)  // Extension Label.
    << uchar(0x0b); // Application Block Size. It says "11 bytes".

  internalStream->write("NETSCAPE2.0", 11); // Extension type, 11 bytes
  p << uchar(0x03)                          // Application block length
    << uchar(0x01)                          // Loop sub-block ID. 1 byte
    << ushort(_repeatCount)                 // Repeat count. 2 bytes.
    << uchar(0x00);                         // Terminator
}

void GifEncoder::writeGraphicControlExtension(int delay, bool isLastFrame) {
  QDataStream p(internalStream);

  p << uchar(0x21)  // Extension Introducer.
    << uchar(0xf9)  // Extension Label.
    << uchar(0x04); // Block size.

  // Packed fields
  uchar bitArray = 0;

  // Lower 3 bits are reserved for future use. Hahahaha. Yeah...

  // Disposal Method

  // Use Inplace if you want to Leave the last frame pixel.
  // GCE_DISPOSAL_NONE = Undefined = 0
  // GCE_DISPOSAL_INPLACE = Leave = 1
  // GCE_DISPOSAL_BACKGROUND = Restore Background = 2
  // GCE_DISPOSAL_RESTORE = Restore Previous = 3

  if (_useFullTransparency) {
    // If full "Transparency" is set:
    // All starting frames as "Restore Background".
    // The last frame as "Leave".

    if (isLastFrame) {
      // Leave.
      bitArray |= 32;
    } else {
      // Restore background.
      bitArray |= 16;
    }
  } else {
    // If "Detect Unchanged Pixels" is set:
    // First frame as "Leave" with no Transparency. IsFirstFrame
    // Following frames as "Undefined" with Transparency.

    // Was TransparentColor.HasValue &&
    if (isFirstFrame) {
      // Leave.
      bitArray |= 32;
    } else {
      // Undefined.
    }
  }

  // User Input Flag. bit 6 = 0

  // Transparent Color Flag, uses tranparency?
  bitArray |=
      ((!isFirstFrame || _useFullTransparency) && colorTableHasTransparency)
          ? 128
          : 0;

  // Write the packed fields.
  p << uchar(bitArray)

    // Calculates the delay, taking into consideration overall rounding. Bug!
    // OrganicTime += delay;
    // delay = (int)Math.Round((OrganicTime > delay ? OrganicTime - AdjustedTime
    // * 10 : delay) / 10.0f, MidpointRounding.AwayFromZero); AdjustedTime +=
    // delay; WriteShort(delay);

    << ushort(qRound(delay / 10.0f))
    << uchar(findTransparentColorIndex()) // Transparency Index.
    << uchar(0);                          // Terminator.
}

void GifEncoder::writeImageDescriptor(QRect rect) {
  QDataStream p(internalStream);

  p << uchar(0x2c)            // Image Separator.
    << ushort(rect.x())       // Position X. 2 bytes.
    << ushort(rect.y())       // Position Y. 2 bytes.
    << ushort(rect.width())   // Width. 2 bytes.
    << ushort(rect.height()); // Height. 2 bytes.

  if (_useGlobalColorTable) {
    // No Local Color Table. Every packed field values are zero.
    p << uchar(0);
    return;
  }

  // Packed fields.
  // Uses local color table?
  // Interlace Flag.
  // Sort Flag.
  // Reserved for future use. Hahahah again.
  uchar bitArray = 5;

  // Size of Local Color Table.
  bitArray |= ((colorTableSize & 7) << 5);

  // Write the packed fields.
  p << uchar(bitArray);
}

void GifEncoder::writeImage() {
  // TODO: Fix the new LZW encoder when ColorTableSize == 7. It's getting
  // corrupted.

  // if (ColorTableSize < 6)
  //{
  //     //New LZW encoder, ColorTableSize from 1 to 8.
  LZWEncoder encoder(indexedPixels, colorTableSize + 1, this);
  encoder.encode(internalStream);
  // }
  // else
  //{
  // Old LZW encoder, only works with ColorTableSize 8.
  // auto encoder = new LZWEncoder(0, 0, indexedPixels, 8);
  // encoder.encode(internalStream);
  //}
}

void GifEncoder::readPixels(QByteArray pixels) {
  if (_quantizationType == ColorQuantizationTypes::Neural) {
    // Neural

    if (globalQuantizer == nullptr || !_useGlobalColorTable) {
      globalQuantizer =
          new NeuralQuantizer(_samplingFactor, _maximumNumberColor);
      globalQuantizer->setMaxColors(_maximumNumberColor);
      globalQuantizer->setTransparentColor(
          !isFirstFrame || _useGlobalColorTable || _useFullTransparency
              ? _transparentColor
              : QColor());
      globalQuantizer->firstPass(pixels);
      colorTable = globalQuantizer->getPalette();
    }

    // Indexes the pixels to the color table.
    indexedPixels = globalQuantizer->secondPass(pixels);

  } else if (_quantizationType == ColorQuantizationTypes::Octree) {
    // Octree

    globalQuantizer = new OctreeQuantizer;
    globalQuantizer->setMaxColors(_maximumNumberColor);
    globalQuantizer->setTransparentColor(
        !isFirstFrame || _useGlobalColorTable || _useFullTransparency
            ? _transparentColor
            : QColor());

    indexedPixels = globalQuantizer->quantize(pixels);
    colorTable = globalQuantizer->colorTable();

  } else if (_quantizationType == ColorQuantizationTypes::MedianCut) {
    // Median cut

    if (globalQuantizer == nullptr || !_useGlobalColorTable) {
      globalQuantizer = new MedianCutQuantizer;
      globalQuantizer->setMaxColors(_maximumNumberColor);
      globalQuantizer->setTransparentColor(
          !isFirstFrame || _useGlobalColorTable || _useFullTransparency
              ? _transparentColor
              : QColor());

      globalQuantizer->firstPass(pixels);
      colorTable = globalQuantizer->getPalette();
    }

    // Indexes the pixels to the color table.
    indexedPixels = globalQuantizer->secondPass(pixels);

  } else if (_quantizationType == ColorQuantizationTypes::Grayscale) {
    // Grayscale

    // This quantizer uses a fixed palette (generated during object
    // instantiation), so most calculations are called one time.
    if (globalQuantizer == nullptr) {
      // Since the color table does not change among frames, it can be stored
      // globally.
      _useGlobalColorTable = true;

      auto transparent =
          !isFirstFrame || _useGlobalColorTable || _useFullTransparency
              ? _transparentColor
              : QColor();

      globalQuantizer =
          new GrayscaleQuantizer(transparent, _maximumNumberColor);
      globalQuantizer->setMaxColors(_maximumNumberColor);
      globalQuantizer->setTransparentColor(transparent);
      colorTable = globalQuantizer->getPalette();
    }
    // Each frame still needs to be quantized.
    indexedPixels = globalQuantizer->secondPass(pixels);

  } else if (_quantizationType == ColorQuantizationTypes::MostUsed) {
    // Most used colors

    if (globalQuantizer == nullptr || !_useGlobalColorTable) {
      globalQuantizer = new MostUsedQuantizer;
      globalQuantizer->setMaxColors(_maximumNumberColor);
      globalQuantizer->setTransparentColor(
          !isFirstFrame || _useGlobalColorTable || _useFullTransparency
              ? _transparentColor
              : QColor());
      globalQuantizer->firstPass(pixels);
      colorTable = globalQuantizer->getPalette();
    }
    // Indexes the pixels to the color table.
    indexedPixels = globalQuantizer->secondPass(pixels);

  } else {
    // Palette

    // This quantizer uses a fixed palette (generated during object
    // instantiation), so it will be only called once.
    if (globalQuantizer == nullptr) {
      // Since the color table does not change among frames, it can be
      // stored globally.
      _useGlobalColorTable = true;

      // TODO: Pass the palette.
      // Default palettes: Windows, etc.
      // User submitted > Presets > Generate palette based on first frame.

      globalQuantizer = new PaletteQuantizer(QVector<QColor>());
      globalQuantizer->setMaxColors(_maximumNumberColor);
      globalQuantizer->setTransparentColor(_transparentColor);
      colorTable = globalQuantizer->getPalette();
    }

    // Each frame still needs to be quantized.
    indexedPixels = globalQuantizer->secondPass(pixels);
  }

  // I need to signal the other method that I'll need transparency.
  colorTableHasTransparency =
      _transparentColor.isValid() && colorTable.contains(_transparentColor);
}

void GifEncoder::writeComment(QString comment) {
  QDataStream p(internalStream);
  p << uchar(0x21) << uchar(0xfe);
  internalStream->write(comment.toLatin1());
}

void GifEncoder::calculateColorTableSize() {
  // Logical Screen Description, Number of Colors, Byte length.
  // 0 = 2 = 6
  // 1 = 4 = 12
  // 2 = 8 = 24
  // 3 = 16 = 48
  // 4 = 32 = 96
  // 5 = 64 = 192
  // 6 = 128 = 384
  // 7 = 256 = 768
  // The inverse calculation is: 2^(N + 1)
  // and x3 for the byte length.

  // If the colorsCount == 1,
  // return zero instead of calculating it, because of the Log(0) call.
  // The "-1" assures that the count stays in range.
  colorTableSize =
      colorTable.count() > 1 ? (int)log2(colorTable.count() - 1) : 0;
}

int GifEncoder::getMaximumColorCount() {
  // 2^(N+1)
  return (int)qPow(2, colorTableSize + 1);
}

int GifEncoder::findTransparentColorIndex() {
  if ((isFirstFrame && !_useFullTransparency) || !colorTableHasTransparency)
    return 0;

  // ReSharper disable once PossibleInvalidOperationException
  auto index = colorTable.indexOf(_transparentColor);
  return index > -1 ? index : 0;
}

} // namespace WingGif
